{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fCV0WMtaFYO6"
      },
      "outputs": [],
      "source": [
        "# Copyright 2018 The TensorFlow GAN Authors. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# =============================================================================="
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RB2DOo-xFoZ3"
      },
      "source": [
        "# TF-GAN on TPUs Tutorial\n",
        "\n",
        "Tutorial authors: joelshor@google.com, westbrook@google.com, tomfeelslucky@google.com"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UUIJxquv6OHk"
      },
      "source": [
        "## Overview\n",
        "\n",
        "\n",
        "This colab walks you through the basics of using TensorFlow GAN (TF-GAN) on Tensor Processing Units (TPUs).\n",
        "\n",
        "This notebook is hosted on GitHub. To view it in its original repository, after opening the notebook, select **File \u003e View on GitHub**.\n",
        "\n",
        "### TPUs\n",
        "\n",
        "TPUs are chips optimized for machine learning training and inference. GAN training uses a lot of compute power, so TPUs expand the range of what we can realistically accomplish with GANs.\n",
        "\n",
        "### CIFAR10 Task\n",
        "\n",
        "In this colab we'll use TPUs to train a GAN for a more difficult task than the MNIST task tackled in the [GANEstimator](https://goto.google.com/tfgan-tutorial) colab. We'll use the [CIFAR10](https://wikipedia.org/wiki/CIFAR-10) dataset to train a model to generate images of airplanes, cars, birds, cats, deer, dogs, frogs, horses, ships, and trucks.\n",
        "\n",
        "The training runs for 50000 steps and completes in roughly 10 minutes on TPUs.\n",
        "\n",
        "### DCGAN Architecture\n",
        "\n",
        "Our model implements the Deep Convolutional Generative Adversarial Network (DCGAN) architecture. [Convolutional Neural Networks](https://developers.google.com/machine-learning/glossary/#convolutional_neural_network) (CNNs) are a common machine learning technique for dealing with images. DCGAN adapts a CNN architecture so that it can function as part of a GAN."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wPCQ0ejbB7eD"
      },
      "source": [
        "## Learning objectives\n",
        "\n",
        "In this Colab, you will learn how to:\n",
        "*   Build a simple DCGAN generator and discriminator\n",
        "*   Use a TPUGANEstimator from TF-GAN to train on CIFAR10 data.\n",
        "*   Run the model forward and see how well it generates images.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C3qZKm2uCKeW"
      },
      "source": [
        "## Instructions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "drBSIcprCM5Q"
      },
      "source": [
        "\u003ch3\u003e  \u0026nbsp;\u0026nbsp;Train on TPU\u0026nbsp;\u0026nbsp; \u003ca href=\"https://cloud.google.com/tpu/\"\u003e\u003cimg valign=\"middle\" src=\"https://raw.githubusercontent.com/GoogleCloudPlatform/tensorflow-without-a-phd/master/tensorflow-rl-pong/images/tpu-hexagon.png\" width=\"50\"\u003e\u003c/a\u003e\u003c/h3\u003e\n",
        "\n",
        "### Steps to run this notebook\n",
        "\n",
        "This notebook should be run in Colaboratory. If you are viewing this from GitHub, follow the GitHub instructions. If you are viewing this from Colaboratory, you should skip to the Colaboratory instructions.\n",
        "\n",
        "#### Steps from GitHub\n",
        "\n",
        "1. Navigate your web brower to the main Colaboratory website: https://colab.research.google.com.\n",
        "1. Click the `GitHub` tab.\n",
        "1. In the field marked `Enter a GitHub URL or search by organization or user`, put in the URL of this notebook in GitHub and click the magnifying glass icon next to it.\n",
        "1. Run the notebook in colaboratory by following the instructions below.\n",
        "\n",
        "#### Steps from Colaboratory\n",
        "\n",
        "1. Create a Cloud Storage bucket for your TensorBoard logs at http://console.cloud.google.com/storage.\n",
        "1. Go to `Runtime \u003e Change runtime type`.\n",
        "1. Click `Hardware accelerator`.\n",
        "1. Select `TPU` and click `Save`.\n",
        "1. Click `Connect` in the upper right corner and select `Connect to hosted runtime`.\n",
        "1. Click Runtime again and select **Runtime \u003e Run All** (Watch out: the initial authentication step for this notebook  requires that you click on `use_tpu` and supply a bucket name as input). You can also run the cells manually with Shift-ENTER.  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u7A6mkRZGRqJ"
      },
      "source": [
        "## Testing out the TPU connection\n",
        "\n",
        "First, you'll need to enable TPUs for the notebook.\n",
        "\n",
        "Navigate to Edit→Notebook Settings, and select TPU from the Hardware Accelerator drop-down (you can also access Notebook Settings via the command palette: cmd/ctrl-shift-P).\n",
        "\n",
        "Next, we'll check that we can connect to the TPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yjFF0NJMHXdm"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import tensorflow.compat.v1 as tf\n",
        "tf.disable_eager_execution()\n",
        "import pprint\n",
        "assert 'COLAB_TPU_ADDR' in os.environ, 'Did you forget to switch to TPU?'\n",
        "tpu_address = 'grpc://' + os.environ['COLAB_TPU_ADDR']\n",
        "\n",
        "with tf.Session(tpu_address) as sess:\n",
        "  devices = sess.list_devices()\n",
        "pprint.pprint(devices)\n",
        "device_is_tpu = [True if 'TPU' in str(x) else False for x in devices]\n",
        "assert True in device_is_tpu, 'Did you forget to switch to TPU?'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MeoJnwf8CT5J"
      },
      "source": [
        "### Authentication"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aC99NXt-AUHJ"
      },
      "source": [
        "To run on Google's free Cloud TPUs, you must set up a [Google Cloud Storage bucket](https://cloud.google.com/storage/) to store logs and checkpoints. Even though this may require entering payment information, running this notebook alone should fall under the [free pricing tier](https://cloud.google.com/storage/pricing) and cost nothing."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "by1-zWuuCWsD"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import os\n",
        "import pprint\n",
        "import re\n",
        "import time\n",
        "import tensorflow.compat.v1 as tf\n",
        "import tensorflow_gcs_config\n",
        "\n",
        "# Google Cloud Storage bucket for Estimator logs and storing\n",
        "# the training dataset.\n",
        "bucket = '' #@param {type:\"string\"}\n",
        "\n",
        "assert bucket, 'Must specify an existing GCS bucket name'\n",
        "print('Using bucket: {}'.format(bucket))\n",
        "\n",
        "model_dir = 'gs://{}/{}'.format(\n",
        "    bucket, time.strftime('tpuestimator-tfgan/%Y-%m-%d-%H-%M-%S'))\n",
        "print('Using model dir: {}'.format(model_dir))\n",
        "\n",
        "assert 'COLAB_TPU_ADDR' in os.environ, 'Missing TPU; did you request a TPU in Notebook Settings?'\n",
        "tpu_address = 'grpc://{}'.format(os.environ['COLAB_TPU_ADDR'])\n",
        "\n",
        "from google.colab import auth\n",
        "auth.authenticate_user()\n",
        "\n",
        "# Upload credentials to TPU.\n",
        "tf.config.experimental_connect_to_host(tpu_address)\n",
        "tensorflow_gcs_config.configure_gcs_from_colab_auth()\n",
        "# Now credentials are set for all future sessions on this TPU."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wiV9VpWKbnLN"
      },
      "source": [
        "### Check imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nXDNWK1BbpJm"
      },
      "outputs": [],
      "source": [
        "# Check that imports for the rest of the file work.\n",
        "import os\n",
        "import tensorflow.compat.v1 as tf\n",
        "!pip install tensorflow-gan\n",
        "import tensorflow_gan as tfgan\n",
        "import tensorflow_datasets as tfds\n",
        "import tensorflow_hub as hub\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "# Allow matplotlib images to render immediately.\n",
        "%matplotlib inline\n",
        "# Disable noisy outputs.\n",
        "tf.logging.set_verbosity(tf.logging.ERROR)\n",
        "tf.autograph.set_verbosity(0, False)\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w7GkdNB6QB7b"
      },
      "source": [
        "## Train and evaluate a GAN model on TPU using TF-GAN.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PoVOsRlPCDP8"
      },
      "source": [
        "### Input pipeline\n",
        "\n",
        "Make the train, eval, and predict input functions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3ADoVivQCGfA"
      },
      "outputs": [],
      "source": [
        "import tensorflow.compat.v1 as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "dataset_dir = 'gs://{}/{}'.format(bucket, 'datasets')\n",
        "\n",
        "def input_fn(mode, params):\n",
        "  assert 'batch_size' in params\n",
        "  assert 'noise_dims' in params\n",
        "  bs = params['batch_size']\n",
        "  nd = params['noise_dims']\n",
        "  split = 'train' if mode == tf.estimator.ModeKeys.TRAIN else 'test'\n",
        "  shuffle = (mode == tf.estimator.ModeKeys.TRAIN)\n",
        "  just_noise = (mode == tf.estimator.ModeKeys.PREDICT)\n",
        "  \n",
        "  noise_ds = (tf.data.Dataset.from_tensors(0)\n",
        "              .map(lambda _: tf.random.normal([bs, nd]))\n",
        "              # If 'predict', just generate one batch.\n",
        "              .repeat(1 if just_noise else None))\n",
        "  \n",
        "  if just_noise:\n",
        "    return noise_ds\n",
        "\n",
        "  def _preprocess(element):\n",
        "    # Map [0, 255] to [-1, 1].\n",
        "    images = (tf.cast(element['image'], tf.float32) - 127.5) / 127.5\n",
        "    return images\n",
        "\n",
        "  images_ds = (tfds.load('cifar10:3.*.*', split=split, data_dir=dataset_dir)\n",
        "               .map(_preprocess, num_parallel_calls=4)\n",
        "               .cache()\n",
        "               .repeat())\n",
        "  if shuffle:\n",
        "    images_ds = images_ds.shuffle(\n",
        "        buffer_size=10000, reshuffle_each_iteration=True)\n",
        "  images_ds = (images_ds.batch(bs, drop_remainder=True)\n",
        "               .prefetch(tf.data.experimental.AUTOTUNE))\n",
        "\n",
        "  return tf.data.Dataset.zip((noise_ds, images_ds))\n",
        "\n",
        "\n",
        "def noise_input_fn(params):\n",
        "  np.random.seed(0)\n",
        "  np_noise = np.random.randn(params['batch_size'], params['noise_dims'])\n",
        "  return tf.data.Dataset.from_tensors(tf.constant(np_noise, dtype=tf.float32))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NDQvDSa8MHGh"
      },
      "source": [
        "Download the data. TensorFlow Datsets will write the data once to your GCS bucket, then reuse it for future calls."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRXkUnr-MIEh"
      },
      "outputs": [],
      "source": [
        "params = {'batch_size': 1, 'noise_dims':1}\n",
        "input_fn(tf.estimator.ModeKeys.EVAL, params)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rW4SlJk0COjP"
      },
      "source": [
        "Sanity check the inputs.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X9vWFknACPvF"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import tensorflow_datasets as tfds\n",
        "import tensorflow_gan as tfgan\n",
        "\n",
        "params = {'batch_size': 80, 'noise_dims':64}\n",
        "ds = input_fn(tf.estimator.ModeKeys.EVAL, params)\n",
        "numpy_imgs = next(iter(tfds.as_numpy(ds)))[1]\n",
        "image_grid = tfgan.eval.python_image_grid(numpy_imgs, grid_shape=(8, 10))\n",
        "\n",
        "def _show_image_grid(image_grid):\n",
        "  plt.axis('off')\n",
        "  plt.imshow((image_grid + 1.0) / 2.0,  # [-1, 1] -\u003e [0, 1]\n",
        "             aspect='auto')\n",
        "  plt.show()\n",
        "_show_image_grid(image_grid)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfmvF-krJhQi"
      },
      "source": [
        "### Neural Net Architecture\n",
        "\n",
        "As usual, our GAN has two separate networks:\n",
        "\n",
        "*  A generator that takes input noise and outputs images\n",
        "*  A discriminator that takes images and outputs a probability of being real\n",
        "\n",
        "We define `discriminator()` and `generator()` builder functions that assemble these networks, along with several helper functions that build pieces of these networks. In the \"Estimator\" section below we pass the `discriminator()` and `generator()` functions to the `TPUGANEstimator`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QV_qFuUwJcOb"
      },
      "outputs": [],
      "source": [
        "def _leaky_relu(x):\n",
        "  return tf.nn.leaky_relu(x, alpha=0.2)\n",
        "\n",
        "\n",
        "def _batch_norm(x, is_training, name):\n",
        "  return tf.layers.batch_normalization(\n",
        "      x, momentum=0.9, epsilon=1e-5, training=is_training, name=name)\n",
        "\n",
        "\n",
        "def _dense(x, channels, name):\n",
        "  return tf.layers.dense(\n",
        "      x, channels,\n",
        "      kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
        "      name=name)\n",
        "\n",
        "\n",
        "def _conv2d(x, filters, kernel_size, stride, name):\n",
        "  return tf.layers.conv2d(\n",
        "      x, filters, [kernel_size, kernel_size],\n",
        "      strides=[stride, stride], padding='same',\n",
        "      kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
        "      name=name)\n",
        "\n",
        "\n",
        "def _deconv2d(x, filters, kernel_size, stride, name):\n",
        "  return tf.layers.conv2d_transpose(\n",
        "      x, filters, [kernel_size, kernel_size],\n",
        "      strides=[stride, stride], padding='same',\n",
        "      kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
        "      name=name)\n",
        "\n",
        "\n",
        "def discriminator(images, unused_conditioning, is_training=True,\n",
        "                  scope='Discriminator'):\n",
        "  \"\"\"Discriminator for CIFAR images.\n",
        "\n",
        "  Args:\n",
        "    images: A Tensor of shape [batch size, width, height, channels], that can be\n",
        "      either real or generated. It is the discriminator's goal to distinguish\n",
        "      between the two.\n",
        "    unused_conditioning: The TFGAN API can help with conditional GANs, which\n",
        "      would require extra `condition` information to both the generator and the\n",
        "      discriminator. Since this example is not conditional, we do not use this\n",
        "      argument.\n",
        "    is_training: If `True`, batch norm uses batch statistics. If `False`, batch\n",
        "      norm uses the exponential moving average collected from population\n",
        "      statistics.\n",
        "    scope: A variable scope or string for the discriminator.\n",
        "\n",
        "  Returns:\n",
        "    A 1D Tensor of shape [batch size] representing the confidence that the\n",
        "    images are real. The output can lie in [-inf, inf], with positive values\n",
        "    indicating high confidence that the images are real.\n",
        "  \"\"\"\n",
        "  with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):\n",
        "    x = _conv2d(images, 64, 5, 2, name='d_conv1')\n",
        "    x = _leaky_relu(x)\n",
        "\n",
        "    x = _conv2d(x, 128, 5, 2, name='d_conv2')\n",
        "    x = _leaky_relu(_batch_norm(x, is_training, name='d_bn2'))\n",
        "\n",
        "    x = _conv2d(x, 256, 5, 2, name='d_conv3')\n",
        "    x = _leaky_relu(_batch_norm(x, is_training, name='d_bn3'))\n",
        "\n",
        "    x = tf.reshape(x, [-1, 4 * 4 * 256])\n",
        "\n",
        "    x = _dense(x, 1, name='d_fc_4')\n",
        "\n",
        "    return x\n",
        "\n",
        "\n",
        "def generator(noise, is_training=True, scope='Generator'):\n",
        "  \"\"\"Generator to produce CIFAR images.\n",
        "\n",
        "  Args:\n",
        "    noise: A 2D Tensor of shape [batch size, noise dim]. Since this example\n",
        "      does not use conditioning, this Tensor represents a noise vector of some\n",
        "      kind that will be reshaped by the generator into CIFAR examples.\n",
        "    is_training: If `True`, batch norm uses batch statistics. If `False`, batch\n",
        "      norm uses the exponential moving average collected from population\n",
        "      statistics.\n",
        "    scope: A variable scope or string for the generator.\n",
        "\n",
        "  Returns:\n",
        "    A single Tensor with a batch of generated CIFAR images.\n",
        "  \"\"\"\n",
        "  with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):\n",
        "    net = _dense(noise, 4096, name='g_fc1')\n",
        "    net = tf.nn.relu(_batch_norm(net, is_training, name='g_bn1'))\n",
        "\n",
        "    net = tf.reshape(net, [-1, 4, 4, 256])\n",
        "\n",
        "    net = _deconv2d(net, 128, 5, 2, name='g_dconv2')\n",
        "    net = tf.nn.relu(_batch_norm(net, is_training, name='g_bn2'))\n",
        "\n",
        "    net = _deconv2d(net, 64, 4, 2, name='g_dconv3')\n",
        "    net = tf.nn.relu(_batch_norm(net, is_training, name='g_bn3'))\n",
        "\n",
        "    net = _deconv2d(net, 3, 4, 2, name='g_dconv4')\n",
        "    net = tf.tanh(net)\n",
        "\n",
        "    return net"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FGjCPvu7DnCT"
      },
      "source": [
        "### Eval Utilities\n",
        "\n",
        "We evaluate the performance of our generator using two standard metrics based on a pretrained classifier: the `Inception Score` and the `Frechet Inception Distance`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "spDJ_oEOELqS"
      },
      "outputs": [],
      "source": [
        "import tensorflow.compat.v1 as tf\n",
        "import tensorflow_gan as tfgan\n",
        "import tensorflow_hub as hub\n",
        "import numpy as np\n",
        "tf.disable_eager_execution()\n",
        "\n",
        "eval_batch_size = 4000 #@param\n",
        "images_per_batch = 2000 #@param\n",
        "\n",
        "def get_real_image_logits(num_images, classifier_model):\n",
        "  \"\"\"Returns an array with logits from real images and a CIFAR classifier.\n",
        "  \n",
        "  We normally want many thousands of examples to run eval. However, we can't fit\n",
        "  inference for all of them in memory at once. Instead, we use TF-GAN eval utils\n",
        "  to more efficiently manage memory.\n",
        "\n",
        "  Args:\n",
        "    num_images: Total number of images to produce logits for.\n",
        "    classifier_model: A Python function that takes images and produces logits.\n",
        "\n",
        "  Returns:\n",
        "    A numpy array of logits of shape close to [num_images, ?].\n",
        "  \"\"\"\n",
        "  ds = input_fn(tf.estimator.ModeKeys.TRAIN, \n",
        "                {'batch_size': images_per_batch, 'noise_dims': 1})\n",
        "  iterator = tf.data.make_one_shot_iterator(ds)\n",
        "\n",
        "  cifar_imgs = iterator.get_next()[1]\n",
        "  real_logits = classifier_model(cifar_imgs)\n",
        "  \n",
        "  with tf.train.MonitoredSession() as sess:\n",
        "    logits = sess.run(real_logits)\n",
        "  assert len(logits.shape) == 2\n",
        "  assert logits.shape[0] == num_images\n",
        "  return logits\n",
        "\n",
        "def init_global_real_logits():\n",
        "  \"\"\"Initialize a global variable with classifier logits for real data.\"\"\"\n",
        "  # We can hold all the real logits in memory at once, since CIFAR10 isn't that\n",
        "  # big. Be sure to calculate it only once.\n",
        "  global real_logits\n",
        "  try:\n",
        "    real_logits is not None\n",
        "  except NameError:\n",
        "    with tf.Graph().as_default():\n",
        "      classifier_model = hub.Module(\"https://tfhub.dev/deepmind/ganeval-cifar10-convnet/1\")\n",
        "      real_logits = get_real_image_logits(\n",
        "          eval_batch_size, classifier_model)\n",
        "  assert real_logits.shape == (eval_batch_size, 10)\n",
        "  \n",
        "def calculate_real_data_classifier_score():\n",
        "  \"\"\"Calculate the classifier score on real data logits.\"\"\"\n",
        "  assert real_logits is not None\n",
        "  classifier_score = tfgan.eval.classifier_score_from_logits(real_logits)\n",
        "  with tf.train.MonitoredSession() as sess:\n",
        "    cscore_real = sess.run(classifier_score)\n",
        "  return cscore_real\n",
        "\n",
        "\n",
        "def get_inception_score_and_fid(est):\n",
        "  \"\"\"Calculate our evaluation metrics.\"\"\"\n",
        "  global real_logits\n",
        "  assert real_logits is not None\n",
        "\n",
        "  tf.reset_default_graph()\n",
        "  # We dont' want to hold all the images and activations at once, so use a\n",
        "  # memory-efficient utility.\n",
        "  def sample_fn():\n",
        "    predictions = np.array([x['generated_data'] for x in est.predict(input_fn)])\n",
        "    assert predictions.shape == (images_per_batch, 32, 32, 3)\n",
        "    return predictions\n",
        "  fake_imgs = tf.concat(\n",
        "      [sample_fn() for _ in range(eval_batch_size // images_per_batch)], axis=0)\n",
        "\n",
        "  classifier_fn = hub.Module(\"https://tfhub.dev/deepmind/ganeval-cifar10-convnet/1\")\n",
        "  fake_logits = classifier_fn(fake_imgs)\n",
        "  fake_logits.shape.assert_is_compatible_with([eval_batch_size, 10])\n",
        "\n",
        "  classifier_score = tfgan.eval.classifier_score_from_logits(fake_logits)\n",
        "  fid = tfgan.eval.frechet_classifier_distance_from_activations(\n",
        "      real_logits, fake_logits)\n",
        "\n",
        "  with tf.train.MonitoredSession() as sess:\n",
        "    cscore_np, fid_np = sess.run([classifier_score, fid])\n",
        "  \n",
        "  return cscore_np, fid_np"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sr4rnauKJlHu"
      },
      "source": [
        "### Estimator\n",
        "\n",
        "TF-GAN's `TPUGANEstimator` is like `GANEstimator`, but it extends TensorFlow's `TPUEstimator` class. `TPUEstimator` handles the details of deploying the network on a TPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iliAVH61a9y4"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import tensorflow.compat.v1 as tf\n",
        "import tensorflow_gan as tfgan\n",
        "tf.disable_eager_execution()\n",
        "\n",
        "noise_dims = 1024 #@param\n",
        "generator_lr = 0.0002  #@param\n",
        "discriminator_lr = 0.0002  #@param\n",
        "train_batch_size = 1024  #@param\n",
        "\n",
        "config = tf.estimator.tpu.RunConfig(\n",
        "    model_dir=model_dir,\n",
        "    master=tpu_address,\n",
        "    tpu_config=tf.estimator.tpu.TPUConfig(iterations_per_loop=images_per_batch))\n",
        "est = tfgan.estimator.TPUGANEstimator(\n",
        "    generator_fn=generator,\n",
        "    discriminator_fn=discriminator,\n",
        "    generator_loss_fn=tfgan.losses.modified_generator_loss,\n",
        "    discriminator_loss_fn=tfgan.losses.modified_discriminator_loss,\n",
        "    generator_optimizer=tf.train.AdamOptimizer(generator_lr, 0.5),\n",
        "    discriminator_optimizer=tf.train.AdamOptimizer(discriminator_lr, 0.5),\n",
        "    joint_train=False,  # train G and D sequentially instead of jointly.\n",
        "    train_batch_size=train_batch_size,\n",
        "    predict_batch_size=images_per_batch,\n",
        "    use_tpu=True,\n",
        "    params={'noise_dims': noise_dims},\n",
        "    config=config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7juQranbJu_3"
      },
      "source": [
        "### Train and Eval Loop\n",
        "\n",
        "Train, eval, and vizualize."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wcJ9dJqeJwwk"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "max_steps = 50000 #@param\n",
        "steps_per_eval = 5000 #@param\n",
        "\n",
        "cur_step = 0\n",
        "start_time = time.time()\n",
        "cscores, fids, steps = [], [], []\n",
        "init_global_real_logits()\n",
        "print('Initialized classifier logits for real data.')\n",
        "classifier_score_real_data = calculate_real_data_classifier_score()\n",
        "print('Calculated classifier score for real data.')\n",
        "while cur_step \u003c max_steps:\n",
        "  # Train for a fixed number of steps.\n",
        "  start_step = cur_step\n",
        "  step_to_stop_at = min(cur_step + steps_per_eval, max_steps)\n",
        "  start = time.time()\n",
        "  est.train(input_fn, max_steps=step_to_stop_at)\n",
        "  end = time.time()\n",
        "  cur_step = step_to_stop_at\n",
        "  \n",
        "  # Print some performance statistics.\n",
        "  steps_taken = step_to_stop_at - start_step\n",
        "  time_taken = end - start\n",
        "  steps_per_sec = steps_taken / time_taken\n",
        "  min_since_start = (time.time() - start_time) / 60.0\n",
        "  print(\"Current step: %i, %.4f steps / sec, time since start: %.1f min\" % (\n",
        "      cur_step, steps_per_sec, min_since_start))\n",
        "  \n",
        "  # Calculate some evaluation metrics.\n",
        "  eval_start_time = time.time()\n",
        "  cscore, fid = get_inception_score_and_fid(est)\n",
        "  eval_time = (time.time() - eval_start_time)\n",
        "  cscores.append(cscore)\n",
        "  fids.append(fid)\n",
        "  steps.append(cur_step)\n",
        "  print(\"Classifier score: %.2f / %.2f, FID: %.1f, \"\n",
        "        \"time to calculate eval: %.2f sec\" % (\n",
        "            cscore, classifier_score_real_data, fid, eval_time))\n",
        "  \n",
        "  # Generate and show some predictions.\n",
        "  predictions = np.array(\n",
        "      [x['generated_data'] for x in est.predict(noise_input_fn)])[:80]\n",
        "  image_grid = tfgan.eval.python_image_grid(predictions, grid_shape=(8, 10))\n",
        "  _show_image_grid(image_grid)\n",
        "\n",
        "# Plot the metrics vs step.\n",
        "plt.title('Frechet distance per step')\n",
        "plt.plot(steps, fids)\n",
        "plt.figure()\n",
        "plt.title('Classifier Score per step')\n",
        "plt.plot(steps, cscores)\n",
        "plt.plot(steps, [classifier_score_real_data] * len(steps))\n",
        "plt.figure()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/brain/python/client:tpu_hw_notebook",
        "kind": "private"
      },
      "name": "TF-GAN on TPUs",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "/piper/depot/tensorflow_gan/examples/colab_notebooks/tfgan_on_tpus.ipynb",
          "timestamp": 1625073496461
        },
        {
          "file_id": "/piper/depot/tensorflow_gan/examples/colab_notebooks/tfgan_on_tpus.ipynb",
          "timestamp": 1625016031564
        },
        {
          "file_id": "/piper/depot/tensorflow_gan/examples/colab_notebooks/tfgan_on_tpus.ipynb",
          "timestamp": 1606129140652
        },
        {
          "file_id": "/piper/depot/tensorflow_gan/examples/colab_notebooks/tfgan_on_tpus.ipynb",
          "timestamp": 1582841684055
        },
        {
          "file_id": "/piper/depot/tensorflow_gan/examples/colab_notebooks/tfgan_on_tpus.ipynb",
          "timestamp": 1569547310685
        },
        {
          "file_id": "/piper/depot/tensorflow_gan/examples/colab_notebooks/tfgan_on_tpus.ipynb",
          "timestamp": 1567688657181
        },
        {
          "file_id": "/piper/depot/tensorflow_gan/examples/colab_notebooks/tfgan_on_tpus.ipynb",
          "timestamp": 1564088453771
        },
        {
          "file_id": "/piper/depot/tensorflow_gan/examples/colab_notebooks/tfgan_on_tpus.ipynb",
          "timestamp": 1560586324901
        },
        {
          "file_id": "/piper/depot/tensorflow_gan/examples/colab_notebooks/tfgan_on_tpus.ipynb",
          "timestamp": 1559825146929
        },
        {
          "file_id": "/piper/depot/tensorflow_gan/examples/colab_notebooks/tfgan_on_tpus.ipynb",
          "timestamp": 1559658404997
        },
        {
          "file_id": "/piper/depot/tensorflow_gan/examples/colab_notebooks/tfgan_on_tpus.ipynb",
          "timestamp": 1559231931573
        },
        {
          "file_id": "0Bz8X96EaC_2-ZW9odlhSOEFXdWs",
          "timestamp": 1493398103910
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 2",
      "name": "python2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
